package simtester.cloud;

import static org.junit.Assert.*;

import java.io.InputStream;
import java.io.OutputStream;
import java.util.Map;
import java.util.concurrent.TimeUnit;

import groovy.mock.interceptor.MockFor
import groovy.mock.interceptor.StubFor
import net.schmizz.sshj.common.Message;
import net.schmizz.sshj.common.SSHException;
import net.schmizz.sshj.common.SSHPacket;
import net.schmizz.sshj.connection.ConnectionException;
import net.schmizz.sshj.connection.channel.direct.PTYMode;
import net.schmizz.sshj.connection.channel.direct.Session
import net.schmizz.sshj.connection.channel.direct.SessionChannel
import net.schmizz.sshj.SSHClient
import net.schmizz.sshj.transport.TransportException;
import net.schmizz.sshj.xfer.scp.SCPFileTransfer
import org.junit.After
import org.junit.Before
import org.junit.Test

class CloudInstanceSSHTests {
    def config      = new ConfigSlurper().parse(new File('conf/SimTesterConfig.groovy').toURI().toURL())
    def testAddress = '111.222.333.444'
    def testUser    = 'ec2-user'

    @Test
    void uploadFileToInstance() {
        // SETUP
        def mockSSHClient = new MockFor(SSHClient)
        def connectAddr
        def authArgs = [:]
        def dummy = new Dummy()
        mockSSHClient.demand.with {
            addHostKeyVerifier { verifier ->  }
            connect { addr ->
                connectAddr = addr
            }
            authPublickey { user, key ->
                authArgs.user    = user
                authArgs.authkey = key
            }
            newSCPFileTransfer { ->
                dummy
            }
            disconnect { ->  }
        }

        // DO TEST
        mockSSHClient.use {
            def ssh = new CloudInstanceSSH(address: testAddress, user: testUser, authkey: config.aws.key)
            ssh.upload(new File('testfile.txt'), to: '/home/ec2-user')
        }
        
        // VERIFY
        assert connectAddr      == testAddress
        assert authArgs.user    == testUser
        assert authArgs.authkey == config.aws.key
    }
    
    @Test
    void uploadFileConnectionFailure() {
        // SETUP
        def mockSSHClient = new MockFor(SSHClient)
        mockSSHClient.demand.with {
            addHostKeyVerifier { verifier ->  }
            connect(12) { addr ->
                throw new ConnectException('Connection timed out')
            }
        }
        
        def testFile = new File('testfile.txt')
        def testFileDest = '/home/ec2-user'
        def ssh = new CloudInstanceSSH(address: testAddress, user: testUser, authkey: config.aws.key, test: true)
        
        def errString
        try {
            mockSSHClient.use {
                ssh.upload(testFile, to: testFileDest)
            }
            fail "Failed to throw exception"
        } catch(CloudException e) {
            errString = e.message
        }
        
        assert errString == "Cannot connect to $testAddress: Connection timed out"
    }
    
    @Test
    void execCommandOnClient() {
        // SETUP
        def mockSSHClient = new MockFor(SSHClient)
        def connectAddr
        def authArgs = [:]
        mockSSHClient.demand.with {
            connect { addr ->
                connectAddr = addr
            }
            authPublickey { user, key ->
                authArgs.user    = user
                authArgs.authkey = key
            }
            startSession {
                def dummy = new Dummy()
                dummy
            }
            disconnect { ->  }
        }
        mockSSHClient.ignore('addHostKeyVerifier')
        
        def execCommand
        def cannedOutput = 'out1'
        CloudInstanceSSH.metaClass.execWithOutput = { args ->
            execCommand = args.cmd
            cannedOutput
        }

        // DO TEST
        def cannedCommand = 'ls /home/ec2-user'
        def output
        mockSSHClient.use {
            def ssh = new CloudInstanceSSH(address: testAddress, user: testUser, authkey: config.aws.key)
            output = ssh.exec(cannedCommand)
        }
        
        // VERIFY
        assert connectAddr      == testAddress
        assert authArgs.user    == testUser
        assert authArgs.authkey == config.aws.key
        assert execCommand      == cannedCommand
        assert output           == cannedOutput
    }
    
    @Test
    void execAsyncDoesNotAllocatePTY() {
        // SETUP
        def stubSSHClient = new StubFor(SSHClient)
        stubSSHClient.demand.with {
            connect { addr ->  }
            authPublickey { user, key ->  }
            startSession {
                def dummy = new Dummy()
                dummy
            }
            disconnect { ->  }
        }
        stubSSHClient.ignore('addHostKeyVerifier')
        
        def mockSession = new MockFor(Dummy)
        def actualExecCmd
        mockSession.demand.with {
            exec { cmd ->
                actualExecCmd = cmd
            }
        }
        mockSession.ignore.with {
            allocateDefaultPTY { ->
                fail 'Should not call allocateDefaultPTY()'
            }
        }
        mockSession.ignore('close')
        
        // DO TEST
        def cannedExecCmd = 'testing'
        stubSSHClient.use {
            mockSession.use {
                def ssh = new CloudInstanceSSH(address: testAddress, user: testUser, authkey: config.aws.key)
                ssh.execAsync(cannedExecCmd)
            }
        }
        
        // VERIFY
        assert actualExecCmd == cannedExecCmd
    }
    
    @Test
    void execWithSudoAllocatesPTY() {
        // SETUP
        def stubSSHClient = new StubFor(SSHClient)
        stubSSHClient.demand.with {
            connect { addr ->  }
            authPublickey { user, key ->  }
            startSession {
                def dummy = new Dummy()
                dummy
            }
            disconnect { ->  }
        }
        stubSSHClient.ignore('addHostKeyVerifier')
        
        def actualExecCmd
        CloudInstanceSSH.metaClass.execWithOutput = { args ->
            'out1'
        }
        
        def mockSession = new MockFor(Dummy)
        mockSession.demand.with {
            allocateDefaultPTY { ->  }
        }
        mockSession.ignore('close')
        
        // DO TEST
        stubSSHClient.use {
            mockSession.use {
                def ssh = new CloudInstanceSSH(address: testAddress, user: testUser, authkey: config.aws.key)
                ssh.exec('sudo testing')
            }
        }
    }
    

    // Dummy class for mocking
    class Dummy {
        void upload(file, dest) { }
        void close() { }
    }
}